import torch
from torch.utils.data import DataLoader, Dataset


class CustomDataset(Dataset):
    def __init__(self, data):
        super().__init__()
        self.data = data
        self.features = data['features']
        self.ground_truths = data['ground_truths']
        self.cls_list = data['cls_list']

    def __getitem__(self, index):
        return self.features[index], self.ground_truths[index], self.cls_list[index]

    def __len__(self):
        return len(self.features)


def get_dataloader(batch_size):
    train_datasets = CustomDataset(torch.load("../VOCdevkit/cat_dog_train_data.pth", weights_only=False))
    train_dataloader = DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
    valid_datasets = CustomDataset(torch.load("../VOCdevkit/cat_dog_valid_data.pth", weights_only=False))
    valid_dataloader = DataLoader(valid_datasets, batch_size=batch_size)

    return train_dataloader, valid_dataloader
